#include <ros/ros.h>
#include "quad_msgs/Trajectory.h"
#include <sensor_msgs/PointCloud2.h>
#include "transform_utils/type_transform.h"
#include "transform_utils/display.h"
#include "trajectory_optimization/trajectory_optimization.h"
#include "trajectory_optimization/bspline_smooth.h"
#include <fstream>
#include <cmath>
#include <rosbag/bag.h>
#include <rosbag/view.h>
#include <pcl_conversions/pcl_conversions.h>

#include <Eigen/Core>
#include <Eigen/Dense>
#include <Eigen/Geometry>

#include <pcl/io/pcd_io.h>
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/kdtree/kdtree_flann.h>
#include <pcl/filters/radius_outlier_removal.h>

using Eigen::MatrixXd;
using Eigen::Vector3d;
using std::cout;
using std::endl;
using std::string;
using std::vector;

ros::Publisher optimized_traj_pub, record_traj_pub;
ros::Publisher cloud_map_pub, cloud_map_filter_pub;
Planner::BsplineSmooth bspline_smooth;
Planner::TrajectoryGenerator traj_gen;
Planner::TrajectoryOptimization traj_opt;

pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_map(new pcl::PointCloud<pcl::PointXYZ>);
pcl::KdTreeFLANN<pcl::PointXYZ>::Ptr kdtree(new pcl::KdTreeFLANN<pcl::PointXYZ>);

vector<MatrixXd> coeffs;
vector<double> time_seq;
int opt_seg_num = 8;

void samplePath(vector<Vector3d> &origin_path)
{
    vector<Vector3d> smooth_path, sample_path;
    bspline_smooth.smoothPath(origin_path, smooth_path);
    bspline_smooth.samplePath(smooth_path, sample_path);
    origin_path = sample_path;
}

void optimizeTraj(const vector<Vector3d> &sample_path)
{
    int opt_times = sample_path.size() / opt_seg_num;
    if (opt_times == 0)
        opt_times = 1;
    // NOTE 这里轨迹阶数先固定为5阶
    coeffs = vector<MatrixXd>(3, MatrixXd::Zero(sample_path.size() - 1, 6));
    time_seq.clear();

    for (int s = 0; s < opt_times; s++)
    {
        int start_index = s * opt_seg_num;
        int end_index = (s + 1) * opt_seg_num;
        if (s == opt_times - 1)
            end_index = sample_path.size() - 1;
        cout << "optimizing trajectory from waypoints index " << start_index
             << " to index " << end_index << endl;

        vector<Vector3d> input_waypoints(sample_path.begin() + start_index,
                                         sample_path.begin() + end_index + 1);
        vector<MatrixXd> part_coeffs;
        vector<double> part_time_seq;
        traj_opt.setWayPoints(input_waypoints);
        // TODO 更好地选取边界条件
        Vector3d startVel(0, 0, 0), startAcc(0, 0, 0), goalVel(0, 0, 0), goalAcc(0, 0, 0);
        // 非起点startVel速度绝对值为1
        if (s > 0)
            startVel = (sample_path[start_index + 1] - sample_path[start_index]).normalized();
        // 非终点处速度绝对值为1
        if (s < opt_times - 1)
            goalVel = (sample_path[end_index + 1] - sample_path[end_index]).normalized();

        traj_opt.setBoundaryCondition(startVel, startAcc, goalVel, goalAcc);
        if (traj_opt.optimize() == -1) // -1为nlopt::FAILURE
        {
            ROS_ERROR("Optimizing trajectory failed");
            return;
        }

        traj_opt.getCoeffs(part_coeffs);
        traj_opt.getTimeSeq(part_time_seq);

        // 输出优化结果
        vector<double> cost_val;
        traj_opt.getCostVal(cost_val);
        cout << "cost_time: " << cost_val[0] << ", cost_smooth: " << cost_val[1] << ", cost_collision: " << cost_val[2]
             << ", cost_dynamic_vel: " << cost_val[3] << ", cost_dynamic_acc: " << cost_val[4] << endl;

        for (int i = 0; i < 3; i++)
            coeffs[i].block(start_index, 0, end_index - start_index, 6) = part_coeffs[i];
        time_seq.insert(time_seq.end(), part_time_seq.begin(), part_time_seq.end());
    }
}

void generateTraj(vector<Vector3d> &waypoints)
{
    traj_gen.setWayPoints(waypoints);
    traj_gen.setBoundaryCondition(Vector3d(0, 0, 0), Vector3d(0, 0, 0), Vector3d(0, 0, 0), Vector3d(0, 0, 0));
    traj_gen.solve(3);
    traj_gen.getCoeffs(coeffs);
    traj_gen.getTimeSeq(time_seq);
}

void safetyCheckForTrajectory()
{
    if (time_seq.empty())
        return;
    int segments = time_seq.size();
    int order = coeffs[0].cols();
    bool flag = true;
    for (size_t s = 0; s < segments; s++)
    {
        double duration = time_seq[s];
        for (double t = 0; t < duration; t += 0.2)
        {
            double x = 0, y = 0, z = 0;
            for (int r = 0; r < order; r++)
            {
                x += coeffs[0](s, r) * pow(t, r);
                y += coeffs[1](s, r) * pow(t, r);
                z += coeffs[2](s, r) * pow(t, r);
            }
            pcl::PointXYZ pt(x, y, z);
            std::vector<int> pointIdxNKNSearch(1);
            std::vector<float> pointNKNSquaredDistance(1);
            if (kdtree->nearestKSearch(pt, 1, pointIdxNKNSearch, pointNKNSquaredDistance) > 0)
            {
                float squaredDis = pointNKNSquaredDistance[0];
                double obsDis = std::sqrt(squaredDis);
                if (obsDis < 1.4)
                {
                    cout << "collision at (" << x << ", " << y << ", " << z << "), dis: " << obsDis << endl;
                    flag = false;
                }
            }
            else
                std::cout << "kdtree search error\n";
        }
    }
    if (flag)
        ROS_INFO("Safety Check Passed");
}

void safetyCheckForWaypoints(const vector<Vector3d> &waypoints)
{
    bool flag = true;
    for (auto w : waypoints)
    {
        pcl::PointXYZ pt(w[0], w[1], w[2]);
        std::vector<int> pointIdxNKNSearch(1);
        std::vector<float> pointNKNSquaredDistance(1);
        if (kdtree->nearestKSearch(pt, 1, pointIdxNKNSearch, pointNKNSquaredDistance) > 0)
        {
            float squaredDis = pointNKNSquaredDistance[0];
            double obsDis = std::sqrt(squaredDis);
            if (obsDis < 1.4)
            {
                cout << "collision at (" << w[0] << ", " << w[1] << ", " << w[2] << "), dis: " << obsDis << endl;
                flag = false;
            }
        }
        else
            std::cout << "kdtree search error\n";
    }
    if (flag)
        ROS_INFO("Safety Check Passed");
    else
        ROS_WARN("Safety Check failed");
}

int main(int argc, char *argv[])
{
    ros::init(argc, argv, "optimize_manual_traj");
    ros::NodeHandle nh;

    bool optimize_flag;
    string manual_bagname, optimize_bagname;
    double init_height;
    double filter_radius;
    int filter_num;
    nh.param("optimize_manual_traj/optimize", optimize_flag, false);
    nh.param("optimize_manual_traj/opt_seg_num", opt_seg_num, 8);
    nh.param("optimize_manual_traj/manual_bagname", manual_bagname, string(""));
    nh.param("optimize_manual_traj/optimize_bagname", optimize_bagname, string(""));
    nh.param("optimize_manual_traj/init_height", init_height, 1.4);
    nh.param("optimize_manual_traj/filter_num", filter_num, 1);
    nh.param("optimize_manual_traj/filter_radius", filter_radius, 0.1);

    optimized_traj_pub = nh.advertise<nav_msgs::Path>("optimized_traj_vis", 1);
    record_traj_pub = nh.advertise<nav_msgs::Path>("record_traj_vis", 1);
    cloud_map_pub = nh.advertise<sensor_msgs::PointCloud2>("cloud_map_vis", 1);
    cloud_map_filter_pub = nh.advertise<sensor_msgs::PointCloud2>("cloud_map_filter_vis", 1);

    vector<nav_msgs::Path::Ptr> traj_data;
    try
    {
        rosbag::Bag bag;
        bag.open(manual_bagname, rosbag::bagmode::Read);
        rosbag::View view(bag, rosbag::TopicQuery("record_trag"));

        for (rosbag::MessageInstance const m : view)
        {
            nav_msgs::Path::Ptr s = m.instantiate<nav_msgs::Path>();
            if (s != nullptr)
                traj_data.push_back(s);
        }
    }
    catch (const std::exception &e)
    {
        std::cerr << e.what() << '\n';
        return 0;
    }

    if (traj_data.size() != 1)
    {
        ROS_ERROR("Global trajectory msgs should obtain 1 msg, but now msg size is %d", traj_data.size());
        return 0;
    }

    string cloud_file;
    pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_map(new pcl::PointCloud<pcl::PointXYZ>);

    nh.param("optimize_manual_traj/cloud_filename", cloud_file, string(""));
    if (pcl::io::loadPCDFile(cloud_file, *cloud_map) == -1)
    {
        ROS_ERROR("Cannot read pcd file");
        return 0;
    }

    sensor_msgs::PointCloud2 cloud_map_msgs;
    pcl::toROSMsg(*cloud_map, cloud_map_msgs);
    cloud_map_msgs.header.frame_id = "world";
    cloud_map_pub.publish(cloud_map_msgs);
    ROS_INFO("cloud map visualized");

    pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_map_filter(new pcl::PointCloud<pcl::PointXYZ>);
    pcl::RadiusOutlierRemoval<pcl::PointXYZ> outrem;
    outrem.setInputCloud(cloud_map);
    outrem.setRadiusSearch(filter_radius);
    outrem.setMinNeighborsInRadius(filter_num);
    // outrem.setKeepOrganized(true);
    outrem.filter(*cloud_map_filter);

    sensor_msgs::PointCloud2 cloud_map_filter_msgs;
    pcl::toROSMsg(*cloud_map_filter, cloud_map_filter_msgs);
    cloud_map_filter_msgs.header.frame_id = "world";
    cloud_map_filter_pub.publish(cloud_map_filter_msgs);

    kdtree->setInputCloud(cloud_map_filter);

    vector<Vector3d> waypoints;
    waypoints.push_back(Vector3d(0, 0, init_height));
    for (size_t i = 10; i < traj_data[0]->poses.size(); i++)
    {
        double x = traj_data[0]->poses[i].pose.position.x;
        double y = traj_data[0]->poses[i].pose.position.y;
        double z = traj_data[0]->poses[i].pose.position.z;
        if (waypoints.size() == 1 && (x * x + y * y) < 2.0)
            continue;
        waypoints.push_back(Vector3d(x, y, z));
    }

    Planner::displayTrajectory(waypoints, record_traj_pub);
    ros::Duration(1.0).sleep();
    Planner::displayTrajectory(waypoints, record_traj_pub);
    ros::Duration(1.0).sleep();

    safetyCheckForWaypoints(waypoints);

    bspline_smooth.setParams(3, 1.5);
    samplePath(waypoints);

    std::cout << "waypoints:\n";
    Planner::printVector(waypoints);

    if (optimize_flag)
    {
        traj_opt.setParams(nh);
        traj_opt.setCloudMap(cloud_map, kdtree);

        optimizeTraj(waypoints);
    }
    else
    {
        traj_gen.setParams(5, 0.5);
        generateTraj(waypoints);
    }

    quad_msgs::Trajectory traj_msgs;
    Planner::coeffsToGlobalTrajMsg(coeffs, time_seq, traj_msgs);
    safetyCheckForTrajectory();

    rosbag::Bag bag;
    bag.open(optimize_bagname, rosbag::bagmode::Write);
    bag.write("global_trajectory", ros::Time::now(), traj_msgs);
    ROS_INFO("Write optimize trajectory to rosbag");

    cout << "cloudMap size: " << cloud_map->size() << ", after filtered: "  << cloud_map_filter->size() << endl;

    ros::Rate rate(1);
    while (ros::ok())
    {
        cloud_map_pub.publish(cloud_map_msgs);
        cloud_map_filter_pub.publish(cloud_map_filter_msgs);
        Planner::displayTrajectory(coeffs, time_seq, optimized_traj_pub);
        rate.sleep();
    }

    return 0;
}
